import numpy as np
import pyLikelihood as pyLike
import BinnedAnalysis as BAn
from UpperLimits import UpperLimits
from LikelihoodState import LikelihoodState

class MyBinnedLikeObj(object):
    """
    This is a binned likelihood analysis tools based on Fermi pyLikelihood package.
    """
    def __init__(self, srcMaps, expCube, binnedExpMap, srcModel, irfs):
        self.irfs = irfs
        self.srcMaps, self.srcModel = srcMaps, srcModel
        self.expCube, self.binnedExpMap = expCube, binnedExpMap

        self.obs = BAn.BinnedObs(srcMaps=self.srcMaps, expCube=self.expCube,
                                 binnedExpMap=self.binnedExpMap, irfs=self.irfs)
        self.fitbit = False

    def setERange(self, Emin=0., Emax=1.e10):
        """
        Use Max(Emin, Emin_default) as minimum energy and Min(Emxa, Emax_default)
        as maximum energy in the fit.
        """
        if Emin >= Emax: raise IOError('Emin(%s) >= Emax(%s)' % (Emin, Emax))

        try:
            self.DRM
        except AttributeError:
            pass
        else:
            defEmin, defEmax = self.DRM.energies.min(), self.DRM.energies.max()
            self.newEmin, self.newEmax = max(defEmin, Emin), min(defEmax, Emax)
            if self.newEmin > defEmin or self.newEmax < defEmax:
                self.DRM.setEnergyRange(self.newEmin, self.newEmax)
                self.fitbit = False
            try:
                print '[INFO] ENERGY RANGE of DRM is [%s:%s]' % (self.DRM.emin, self.DRM.emax)
            except AttributeError:
                print '[INFO] ENERGY RANGE of DRM is [%s:%s]' % (self.DRM.energies.min(), self.DRM.energies.max())

        try:
            self.MIN
        except AttributeError:
            pass
        else:
            defEmin, defEmax = self.MIN.energies.min(), self.MIN.energies.max()
            self.newEmin, self.newEmax = max(defEmin, Emin), min(defEmax, Emax)
            if self.newEmin > defEmin or self.newEmax < defEmax:
                self.MIN.setEnergyRange(self.newEmin, self.newEmax)
                self.fitbit = False
            try:
                print '[INFO] ENERGY RANGE of MIN is [%s:%s]' % (self.MIN.emin, self.MIN.emax)
            except AttributeError:
                print '[INFO] ENERGY RANGE of MIN is [%s:%s]' % (self.MIN.energies.min(), self.MIN.energies.max())

    def saveEdisp(self):
        """
        Save the energy dispersion settings in first pyLikelihood Object. You can restore
        the edisp settings to the second pyLikelihood Object by using restoreEdisp method.
        IT MAY BE MEANINGLESS?! Anyway, I keep it.
        """
        try:
            self.DRM
        except AttributeError:
            raise RuntimeError('Please run initDRM first!')
        else:
            if self.irfs.startswith('P8'):
                self.eDispSrcNames = list(self.DRM.model.srcNames)
                self.eDispFlags = []
                print '[INFO] Following are sources taken EDISP into consideration:'
                for src in self.eDispSrcNames:
                    if self.DRM[src].src.use_edisp(): print src+';',
                    self.eDispFlags.append(self.DRM[src].src.use_edisp())
                self.isRestoreEdisp = False # a flag show whether edisp of self.MIN is restored
                print
            else:
                print 'The irfs %s may not support edisp, skip!' % self.irfs

    def restoreEdisp(self, likeObj='MIN'):
        """
        Restore the edisp settings of the first pyLikelihood Object to the second one, you need
        to run saveEdisp before using this method!
        IT MAY BE MEANINGLESS?! Anyway, I keep it.
        """
        hasALTFIT = 'ALTFIT' in dir(self)
        hasMIN = 'MIN' in dir(self)
        if not (hasMIN or hasALTFIT):
            raise RuntimeError('Initiate ALTFIT or MIN first!')

        if not 'isRestoreEdisp' in dir(self):
            print '[ERROR] The irfs you use do not support edisp or you have not run saveEdisp before!'
        elif not self.isRestoreEdisp:
            print '[INFO] Try to recover EDISP settings ...'
            if hasMIN and (likeObj.lower() == 'min'):
                print '[INFO] MIN:'
                nchange = 0
                for isrc, src in enumerate(self.eDispSrcNames):
                    flag_before = self.MIN[src].src.use_edisp()
                    flag_after = self.eDispFlags[isrc]
                    if flag_before ^ flag_after:
                        print '[%s] %s -> %s' % (src, flag_before, flag_after)
                        self.MIN[src].src.set_edisp_flag(flag_after)
                        nchange += 1
                print '[INFO] The edisp setting of %s sources in MIN are recovered!' % nchange

                self.isRestoreEdisp = True

            if hasALTFIT and (likeObj.lower() == 'altfit'):
                print '[INFO] ALTFIT:'
                nchange = 0
                for isrc, src in enumerate(self.eDispSrcNames):
                    flag_before = self.ALTFIT[src].src.use_edisp()
                    flag_after = self.eDispFlags[isrc]
                    if flag_before ^ flag_after:
                        print '[%s] %s -> %s' % (src, flag_before, flag_after)
                        self.ALTFIT[src].src.set_edisp_flag(flag_after)
                        nchange += 1
                print '[INFO] The edisp setting of %s sources in ALTFIT are recovered!' % nchange
        else:
            print '[INFO] YOU HAVE ALREADY RECOVER EDISP SETTINGS, skip!'


    def freezeSource(self, srcName, freezeAll=True, disableEdisp=False, likeObj='DRM'):
        """
        This is a method to freeze sources. Set freezeAll=True to freeze all the
        parameters of the source. Set disableEdisp=True to disable EDISP for the
        sources of whom all the paramters are fixed.
        """
      # if 'iem' in srcName.lower() or 'iso' in srcName.lower():
      #     return False

        if likeObj.lower() == 'min' and 'MIN' in dir(self):
            likeToFreeze = self.MIN
        elif 'DRM' in dir(self):
            likeToFreeze = self.DRM
        else:
            raise RuntimeError('Please run init first!')

        try:
            if likeToFreeze[srcName].src.fixedSpectrum():
                return False
        except AttributeError:
            return False

        hasFreezed = False
        for src_par in likeToFreeze[srcName].funcs['Spectrum'].paramNames:
            if (src_par == 'Prefactor' or src_par == 'norm' or
                src_par == 'Integral' or src_par == 'Value' or
                src_par == 'Normalization'):
                norm_par = src_par
            elif likeToFreeze[srcName].funcs['Spectrum'].params[src_par].isFree():
                src_index = likeToFreeze.par_index(srcName, src_par)
                likeToFreeze.freeze(src_index)
                hasFreezed = True

        if freezeAll or (not hasFreezed):
            print '[INFO] All the parameters of %s are freezed' % srcName
            src_index = likeToFreeze.par_index(srcName, norm_par)
            likeToFreeze.freeze(src_index)

          # if disableEdisp:
          #     try:
          #         likeToFreeze[srcName].src.set_edisp_flag(False)
          #     except AttributeError:
          #         pass
        else:
            print '[INFO] The spectrum of %s is freezed' % srcName

        if disableEdisp:
            try:
                likeToFreeze[srcName].src.set_edisp_flag(False)
            except AttributeError:
                pass
        return True

    def unLoadSource(self, srcName):
        """
        This function removes a source from the model and stores it so that
        you can use it later if you would like.  This is useful if you are
        working on an upper limit and need to get a fit to work before you can
        calculate the upper limit.
        """
        self.saved_src = self.MIN.deleteSource(srcName)
        print '[INFO] %s is removed in self.MIN' % srcName

    def reLoadSource(self):
        """
        This function puts the last source removed by the unLoadSource function
        back into the model.
        """
        try:
            self.MIN.addSource(self.saved_src)
            print '[INFO] Reload the source %s.' % self.saved_src.getName()
        except AttributeError:
            raise RuntimeError("Saved Source does not exist.\
                                Make sure that you have run the unLoadSource function.")

    def freezeWeak(self, centerSrc, tsLimit1=2., tsLimit2=25., distLimit=0.):
        """
        The sources with ts-value < tsLimit1 and distance > distLimit will be totally fixed, between
        tsLimit1 and tsLimit2 spectrum will be fixed. Set disableEdisp=True will set apply_edisp=False
        for totally fixed sources. If there is some source fixed, you must run the fit again!
        """
        try:
            self.currentLike
        except AttributeError:
            raise RuntimeError("MIN object does not exist.\
                  Create it first with the initMIN function and then fit it with the fitMIN function.")

        try:
            self.currentLike[centerSrc].src.getName()
        except AttributeError:
            raise IOError(centerSrc+" is not in the model.  Either pass a valid "+
                          " sourcename to this function or modify your model "+
                          "and/or config file to indicate which source you are considering.")
        freezeAll = []
        freezeSpec = []
        distance = 0.1
        for name in self.currentLike.sourceNames():
            if (self.currentLike[name].src.fixedSpectrum()) or (name == centerSrc) or\
                ('iso_' in name.lower()) or ('iem_' in name.lower()):
                continue
            ts = self.currentLike.Ts(name)
            if self.currentLike[name].src.getType() == 'Point' and self.currentLike[centerSrc].src.getType() == 'Point':
                distance = self.currentLike._separation(self.currentLike[centerSrc].src, self.currentLike[name].src)
                print '[%s] TS=%.3f, distance=%.3f' % (name, ts, distance)
            else:
                print '[%s] TS=%.3f' % (name, ts)

            if ts < tsLimit1 and distance > distLimit:
                freezeAll.append(name)
            elif ts < tsLimit2 and distance > distLimit:
                freezeSpec.append(name)

        if len(freezeSpec) + len(freezeAll) == 0:
            print '[INFO] No weak sources are freezed! You can set larger TS limit (>%s)' % tsLimit2
            return False

        self.currentPristine.restore()
        if 'DRM' in dir(self) and self.currentPristine.like is self.DRM:
            for src in freezeAll:
                self.freezeSource(src, True, True, 'DRM')
            for src in freezeSpec:
                self.freezeSource(src, False, True, 'DRM')
            self.currentPristine = LikelihoodState(self.DRM)
        elif 'MIN' in dir(self) and self.currentPristine.like is self.MIN:
            for src in freezeAll:
                self.freezeSource(src, True, True, 'MIN')
            for src in freezeSpec:
                self.freezeSource(src, False, True, 'MIN')
            self.currentPristine = LikelihoodState(self.MIN)
        else:
            raise RuntimeError('currentPristine.like is neighter MIN nor DRM! DEBUG is required!')
        print '[INFO] THE MODEL IS RESTORED TO THE INITIAL ONE, PLEASE RUN fitDRM AGAIN!'
        return True

    def removeWeak(self, centerSrc, tslimit=0, distlimit=0, RemoveFree=False, RemoveFixed=False):
        """
        This function has two main uses: it will print out details
        on all of the sources in your model and it will remove sources
        according to different requirements.  If you just want to
        print out details, execute it this way:

        <obj>.removeWeak(<my_source>)

        Where <obj> is the quickLike object you're using here and
        <my_source> is the name of your source of interest.  You can
        then remove some of these sources from the model if you like.
        For example, if you want to remove all of the fixed sources
        with TS values less than 1, execute it this way:

        <obj>.removeWeak(<my_source>,tslimit=1,RemoveFixed=True)

        You can mix and match any of the options.  You could remove
        all sources (fixed and free) that are below a TS value of 3
        and are 10 degrees from your source of interest by executing:

        <obj>.removeWeak(<my_source>,tslimit=3,distlimit=10,RemoveFree=True,RemoveFixed=True)
        """
        try:
            self.MIN
        except AttributeError:
            raise RuntimeError("MIN object does not exist.\
                  Create it first with the initMIN function and then fit it with the fitMIN function.")

        if not self.fitbit:
            print "[WARN] Fit isn't current, these values might not be correct.  Run fitMIN first."

        try:
            self.MIN.model[centerSrc].src.getName()
        except AttributeError:
            raise IOError(centerSrc+" is not in the model.  Either pass a valid "+
                          " sourcename to this function or modify your model "+
                          "and/or config file to indicate which source you are considering.")

        print '[INFO] DO NOT USE freezeWeak AND removeWeak AT THE SAME TIME!!! OR ERROR WILL OCCUR!'

        for name in self.MIN.sourceNames():
            if self.MIN[name].src.fixedSpectrum() and (not RemoveFixed): continue
            remove = False
            distance = 0
            sourceTS = self.MIN.Ts(name)
            if self.MIN[name].src.getType() == 'Point' and self.MIN[centerSrc].src.getType() == 'Point':
                distance = self.MIN._separation(self.MIN[centerSrc].src, self.MIN[name].src)
            if self.MIN[name].src.fixedSpectrum():
                indexFree = "Fixed"
                if (sourceTS < tslimit) and (distance > distlimit) and RemoveFixed:
                    remove = True
            else:
                indexFree = "Free"
                if (sourceTS < tslimit) and (distance > distlimit) and RemoveFree:
                    remove = True

            if remove:
                print "[INFO] Removing %s, TS: %s, Frozen?: %s, Distance: %s" % (name, sourceTS, indexFree, distance)
                self.MIN.deleteSource(name)
            else:
                print "[INFO] Retaining %s, TS: %s, Frozen?: %s, Distance: %s" % (name, sourceTS, indexFree, distance)

    def calcUpper(self, srcName, emin=None, emax=None):
        """
        Calculates an upper limit for a source in your model.
        """
        self.ul = UpperLimits(self.MIN)
        if isinstance(emin, float) or isinstance(emin, int):
            _emin = float(emin)
        else:
            try:
                _emin = self.newEmin
            except AttributeError:
                _emin = self.MIN.energies.min()

        if isinstance(emax, float) or isinstance(emax, int):
            _emax = float(emax)
        else:
            try:
                _emax = self.newEmax
            except AttributeError:
                _emax = self.MIN.energies.max()

        print 'Calculating 0.95 Upper Limit flux between [%.3f, %.3f]...' % (_emin, _emax)
        self.ul[srcName].compute(emin=_emin, emax=_emax)
        print 'The upper limit of %s is %s' % (srcName, self.ul[srcName].results)
        return self.ul[srcName].results[0].value, _emin, _emax

    def pokeSource(self, srcName, pname=None):
        """
        This function pokes a paramter of a source to a value that is 10%
        of what it was.  This is a useful function to use when you are trying
        to get convergence on a fit.  Many times, a fit won't converge because
        the initial model is too close to the final answer (ie. the minimizer
        does not have enough flexibility to accurately calculate a correlation
        matrix).  In this case, run this funciton on one of the stronger
        sources in your model and redo the fit.  This function is also useful
        to determine how robust your fit is.
        """
        if pname is None:
            if 'Prefactor' in self.currentLike[srcName].funcs['Spectrum'].paramNames:
                paramName = 'Prefactor'
            elif 'norm' in self.currentLike[srcName].funcs['Spectrum'].paramNames:
                paramName = 'norm'
            elif 'Integral' in self.currentLike[srcName].funcs['Spectrum'].paramNames:
                paramName = 'Integral'
            elif 'Value' in self.currentLike[srcName].funcs['Spectrum'].paramNames:
                paramName = 'Value'
            elif 'Normalization' in self.currentLike[srcName].funcs['Spectrum'].paramNames:
                paramName = 'Normalization'
            else:
                raise RuntimeError('No invalid paramName!')
        else:
            if pname in self.currentLike[srcName].funcs['Spectrum'].paramNames:
                paramName = pname
            else:
                raise RuntimeError('Invalid paramName!')
        previousValue = self.currentLike[srcName].funcs['Spectrum'].getParam(paramName).value()
        self.currentLike[srcName].funcs['Spectrum'].func.setParam(paramName, 0.1*previousValue)
        print ("[INFO] Resetting the %s of %s from %.2f to %.2f" %
               (paramName, srcName, previousValue, 0.1*previousValue))

    def pokeBrightest(self):
        print '[INFO] Check the ts value and find the bright source to poke:'
        brightest, TSmax = None, 0.
        for name in self.currentLike.sourceNames():
            if self.currentLike[name].src.fixedSpectrum(): continue
            if ('iso_' in name.lower()) or ('iem_' in name.lower()): continue
            sourceTS = self.currentLike.Ts(name)
            print '[%s] TS=%.3f' % (name, sourceTS)
            if sourceTS > TSmax:
                brightest, TSmax = name, sourceTS

        if brightest is None:
            raise RuntimeError('The fit fails! No reliable information is available.')

        self.pokeSource(brightest)

    def writeResults(self, outname='results_pyLike.dat', centerSrc=None, TSLimit=None, ULemin=None, ULemax=None):
        """
        Write the fit results in a dict-like fashion.
        Codes belows are adapted from the one of Y-F, Liang
        """
        import pprint

        try:
            self.currentLike
        except AttributeError:
            raise RuntimeError('Please run init first!')

        try:
            Emin, Emax = self.newEmin, self.newEmax
        except AttributeError:
            Emin, Emax = self.currentLike.energies.min(), self.currentLike.energies.max()

        counts = 0
        dicttot = {"logLikelihood" : self.currentLike.logLike.value()}

        if 'MINobj' in dir(self):
            try:
                dicttot['zEDM'] = self.MINobj.getDistance()
                dicttot['zReturnCode'] = self.MINobj.getRetCode()
                dicttot['zFitQuality'] = self.MINobj.getQuality()
            except AttributeError:
                pass

        for src in self.currentLike.sourceNames():
            if self.currentLike[src].src.fixedSpectrum():
                counts += self.currentLike.NpredValue(src)
                continue

            dict_ = {}
            npred = self.currentLike.NpredValue(src)
            ts = self.currentLike.Ts(src)
            flux_ = self.currentLike.flux(src, emin=Emin, emax=Emax, energyFlux=False)
            fluxErr_ = self.currentLike.fluxError(src, emin=Emin, emax=Emax, energyFlux=False)
            counts += npred

            dict_['TS value'] = str(ts)
            dict_['Flux'] = "%.5e +/- %.5e" %(flux_, fluxErr_)
            dict_['NPred'] = str(npred)

            if (centerSrc is None) or (not centerSrc in self.currentLike.sourceNames()):
                print '[%s] TS=%.3f' % (src, ts)
            elif (self.currentLike[centerSrc].src.getType() != 'Point') or\
                 (self.currentLike[src].src.getType() != 'Point'):
                if src == centerSrc:
                    print '[%s] TS=%.3f [CENTER SOURCE]' % (src, ts)
                else:
                    print '[%s] TS=%.3f' % (src, ts)
            else:
                distance = self.currentLike._separation(self.currentLike[centerSrc].src, self.currentLike[src].src)
                dict_['distance'] = '%.3f' % distance
                if src == centerSrc:
                    print '[%s] TS=%.3f, distance=%.3f [CENTER SOURCE]' % (src, ts, distance)
                else:
                    print '[%s] TS=%.3f, distance=%.3f' % (src, ts, distance)

            for pname in self.currentLike[src].funcs['Spectrum'].paramNames:
                error_ = self.currentLike[src].funcs['Spectrum'].params[pname].error()
                value_ = self.currentLike[src].funcs['Spectrum'].params[pname].value()
                dict_[pname] = "%f +/- %f" %(value_, error_)
            dicttot[src] = dict_

        if (not centerSrc is None) and (centerSrc in self.currentLike.sourceNames()) and\
           (not TSLimit is None):
            if float(dicttot[centerSrc]['TS value']) < TSLimit:
                flux_ul, _emin, _emax = self.calcUpper(centerSrc, emin=ULemin, emax=ULemax)
                dicttot[centerSrc]['Upper Limit'] = str(flux_ul)
                dicttot[centerSrc]['ULERange'] = str((_emin, _emax))

        f = open(outname, 'w')
        pprint.pprint(dicttot, stream=f)
        f.close()
        print ("[INFO] Model Cnts / Obs Cnts: %s / %s = 1.%+.5f" %
               (counts, self.currentLike.total_nobs(), counts/self.currentLike.total_nobs()-1.))

    def initDRM(self, srcModel=None, tol=0.001):
        """
        Initializes the DRM optimizer. To get good convergence, you need to fit the
        model using this method first along with fitDRM, and then run minuit.
        """
        if srcModel is None: srcModel = self.srcModel
        self.DRM = BAn.BinnedAnalysis(self.obs, srcModel, optimizer="DRMNFB")
        self.DRM.tol = float(tol)
        self.currentLike = self.DRM
        self.DRMpristine = LikelihoodState(self.DRM)
        self.currentPristine = self.DRMpristine
        print '[INFO] Something about DRM:'
        print self.DRM

    def initAltFit(self, srcModel=None, tol=0.001):
        """
        If DRMNFB fails, the class will initiate this alternative pyLikelihood Object
        and run another fit.
        """
        if srcModel is None: srcModel = self.srcModel
        self.ALTFIT = BAn.BinnedAnalysis(self.obs, srcModel, optimizer='MINUIT')
        self.ALTFIT.tol = float(tol)
        if 'newEmin' in dir(self):
            defEmin, defEmax = self.ALTFIT.energies.min(), self.ALTFIT.energies.max()
            if self.newEmin > defEmin or self.newEmax < defEmax:
                self.ALTFIT.setEnergyRange(self.newEmin, self.newEmax)
            try:
                print '[INFO] ENERGY RANGE of ALTFIT is [%s:%s]' % (self.ALTFIT.emin, self.ALTFIT.emax)
            except AttributeError:
                print '[INFO] ENERGY RANGE of ALTFIT is [%s:%s]' % (self.ALTFIT.energies.min(), self.ALTFIT.energies.max())
        self.ALTFITobj = pyLike.Minuit(self.ALTFIT.logLike)
        self.currentLike = self.ALTFIT
        print '[INFO] Something about ALTFIT:'
        print self.ALTFIT

    def initMIN(self, opt='Minuit', tol=1.e-5, useBadFit=False):
        """
        Initiallizes a New Minuit optimizer to use as a backup to
        the DRM optimizer.  This is usually run after you have
        initially run fitDRM and created a model_pass1.xml
        model file which is used a seed for the New Minuit optimizer.
        You can skip the DRM process if you like but you need to have
        the proper model file (model_pass1.xml) present in the
        working directory.  If you want to use the non convergant fit
        from fitDRM, set useBadFit to True.
        """
        self.useBadFit = useBadFit
        try:
            if self.badfit and self.useBadFit:
                modelFile = self.modelPass1BAD
            else:
                modelFile = self.modelPass1
        except AttributeError:
            modelFile = self.srcModel

        print '[INFO] PASS2 MODEL: ' + modelFile
        if opt.lower() == 'minuit':
            self.MIN = BAn.BinnedAnalysis(self.obs, modelFile, optimizer='MINUIT')
            self.MINobj = pyLike.Minuit(self.MIN.logLike)
        else:
            self.MIN = BAn.BinnedAnalysis(self.obs, modelFile, optimizer='NewMinuit')
            self.MINobj = pyLike.NewMinuit(self.MIN.logLike)

        if 'newEmin' in dir(self):
            defEmin, defEmax = self.MIN.energies.min(), self.MIN.energies.max()
            if self.newEmin > defEmin or self.newEmax < defEmax:
                self.MIN.setEnergyRange(self.newEmin, self.newEmax)
            try:
                print '[INFO] ENERGY RANGE of MIN is [%s:%s]' % (self.MIN.emin, self.MIN.emax)
            except AttributeError:
                print '[INFO] ENERGY RANGE of MIN is [%s:%s]' % (self.MIN.energies.min(), self.MIN.energies.max())

        self.MIN.tol = float(tol)
        self.MINpristine = LikelihoodState(self.MIN)
        self.currentLike = self.MIN
        self.currentPristine = self.MINpristine
        if 'isRestoreEdisp' in dir(self): self.isRestoreEdisp = False
        print '[INFO] Something about MIN:'
        print self.MIN

    def fitDRM(self, modelOut='model_pass1.xml', pokeSrc=None):
        """
        Performs a DRM inital fit on your data using the
        model_pass1.xml model file.  It tries an intial fit and
        if that fails, tries a tighter tolerance.  If that fails, it
        tries a looser tolerance.  If that fails, it tries to do this
        initial fit with the MINUIT optimizer.  If that fails, this
        function bails.  If the fit converges, it saves the results to
        <basename>_likeDRM.xml which will be used in the NewMinuit
        fit.  If no fit is found, it will save the results to
        <basename>_badDRMFit.xml.  You can use this in the NewMinuit fit
        if you use the useBadFit option in initMIN.  You need to have run
        initDRM before you run this function.
        """
        try:
            self.DRM
        except AttributeError:
            raise RuntimeError("DRM object does not exist.  Create it first with the initDRM function.")

        tol1 = self.DRM.tol
        altfit = False
        loglike_old = self.DRM.logLike.value()
        try:
            self.DRM.fit()
            if self.DRM.logLike.value() < loglike_old:
                raise RuntimeError('logLike[%s] < logLike_old[%s]' % (self.DRM.logLike.value(), loglike_old))
        except RuntimeError as e:
            print "[ERROR] Initial DRM Fit Failed\n%sTotal TS: %s" % (e, self.DRM.logLike.value())
            try:
                print "[INFO] Trying tighter tolerance (DRMtol*0.1)"
                self.DRM.tol = tol1 * 0.1
                if not pokeSrc is None: self.pokeSource(pokeSrc)
                self.DRM.fit()
                if self.DRM.logLike.value() < loglike_old:
                    raise RuntimeError('logLike[%s] < logLike_old[%s]' % (self.DRM.logLike.value(), loglike_old))
            except RuntimeError as e:
                print "[ERROR] Second DRM Fit Failed\n%sTotal TS: %s" % (e, self.DRM.logLike.value())
                try:
                    print "[INFO] Trying looser tolerance (DRMtol*10.)"
                    self.DRM.tol = tol1 * 10.
                    if not pokeSrc is None: self.pokeSource(pokeSrc)
                    self.DRM.fit()
                    if self.DRM.logLike.value() < loglike_old:
                        raise RuntimeError('logLike[%s] < logLike_old[%s]' % (self.DRM.logLike.value(), loglike_old))
                except RuntimeError as e:
                    print "[ERROR] Third DRM Fit Failed\n%sTotal TS: %s" % (e, self.DRM.logLike.value())
                    self.saveEdisp()
                    self.currentPristine.restore()
                    altSrcModel = '/tmp/'+'.'.join(modelOut.split('/')[-1].split('.')[:-1]) + '_alt_%05i.xml' %np.random.randint(0, 1e5)
                  # altSrcModel = '.'.join(self.srcModel.split('.')[:-1]) + '_alt.xml'
                    self.DRM.logLike.writeXml(altSrcModel)
                    print '[INFO] Save current model as ' + altSrcModel
                    print "[INFO] Trying MINUIT algorithm ..."
                    self.initAltFit(srcModel=altSrcModel, tol=tol1*20.)
                    self.restoreEdisp('altfit')
                    self.ALTFIT.fit(covar=True, optObject=self.ALTFITobj)
                    print '[INFO] Fit Quality:  ', self.ALTFITobj.getQuality()
                    print '[INFO] Return Code:  ', self.ALTFITobj.getRetCode()
                    print '[INFO] Fit Distance: ', self.ALTFITobj.getDistance()
                    if self.ALTFITobj.getQuality() < 3:
                        print "[ERROR] Alternative fit algorithm failed, bailing"
                        self.badfit = True
                        self.modelPass1BAD = '.'.join(modelOut.split('.')[:-1]) + '_BAD.xml'
                        self.ALTFIT.logLike.writeXml(self.modelPass1BAD)
                        print "[INFO] Saved ALTFIT as " + self.modelPass1BAD
                        return False
                    else:
                        altfit = True
                    if self.ALTFIT.logLike.value() < loglike_old:
                        raise RuntimeError('Some fatal errors may occur in fit algorithm? Check the log!')

        self.badfit = False
        self.modelPass1 = modelOut

        if altfit:
            print "[INFO] ALTFIT Fit Finished.  Total TS: "+str(self.ALTFIT.logLike.value())
            self.ALTFIT.logLike.writeXml(self.modelPass1)
            print "[INFO] Saved ALTFIT as "+self.modelPass1
        else:
            print "[INFO] DRM Fit Finished.  Total TS: "+str(self.DRM.logLike.value())
            self.DRM.logLike.writeXml(self.modelPass1)
            print "[INFO] Saved DRM as "+self.modelPass1
        return True

    def fitMIN(self, modelOut='model_pass2.xml'):
        """
        Does a Minuit fit on your data based on the model output by the fitDRM function.
        You need to have run initMIN before running this function.
        """
        try:
            self.MIN
        except AttributeError:
            raise RuntimeError("MIN object does not exist.  Create it first with the initMIN function.")

        try:
            self.MINobj.getQuality()
        except AttributeError:
            raise RuntimeError('MIN should be a Minuit PyLikelihood Object')

        if ('isRestoreEdisp' in dir(self)) and (not self.isRestoreEdisp): self.restoreEdisp()
        self.MIN.fit(covar=True, optObject=self.MINobj)

        print '[INFO] MINUIT Fit Finished.'
        print '[INFO] Total TS: ', self.MIN.logLike.value()
        print '[INFO] Fit Quality:  ', self.MINobj.getQuality()
        print '[INFO] Return Code:  ', self.MINobj.getRetCode()
        print '[INFO] Fit Distance: ', self.MINobj.getDistance()

        self.fitbit = True
        if self.MINobj.getQuality() < 3:
            print "[ERROR] MINUIT DID NOT CONVERGE!!!"
            self.modelPass2BAD = '.'.join(modelOut.split('.')[:-1]) + '_BAD.xml'
            self.MIN.logLike.writeXml(self.modelPass2BAD)
            print "[INFO] Saved MIN as "+self.modelPass2BAD
            return False
        else:
            self.modelPass2 = modelOut
            self.MIN.logLike.writeXml(self.modelPass2)
            print "[INFO] Saved MIN as "+self.modelPass2
            return True

    def fitMIN_NewMinuit(self, modelOut='model_pass2.xml'):
        """
        Does a New Minuit fit on your data based on the model
        output by the fitDRM function.  You need to have run initMIN
        before running this function.  Saves the results to
        <basename>_likeMIN.xml if there is convergence.  If
        convergence is not found, saves the results to
        <basename>_badMINFit.xml.
        [ATTENTION] NewMinuit seems to be worse than Minuit
        """
        try:
            self.MIN
        except AttributeError:
            raise RuntimeError("MIN object does not exist.  Create it first with the initMIN function.")

        try:
            self.MINobj.getQuality()
        except AttributeError:
            pass
        else:
            raise RuntimeError('MIN should be a NewMinuit PyLikelihood Object')

        tol2 = self.MIN.tol
        if ('isRestoreEdisp' in dir(self)) and (not self.isRestoreEdisp): self.restoreEdisp()
        try:
            self.MIN.fit(covar=True, optObject=self.MINobj)
        except RuntimeError as e:
            print "[ERROR] Initial MIN Fit Failed\n%sTotal TS: %s" % (e, self.MIN.logLike.value())
            try:
                print "[INFO] Trying tighter tolerance (MINtol*0.1)"
                self.MIN.fit(covar=True, optObject=self.MINobj, tol=tol2*0.1)
            except RuntimeError as e:
                print "[ERROR] Second MIN Fit Failed\n%sTotal TS: %s" % (e, self.MIN.logLike.value())
                try:
                    print "[INFO] Trying looser tolerance (MINtol*10.)"
                    self.MIN.fit(covar=True, optObject=self.MINobj, tol=tol2*10.)
                except RuntimeError as e:
                    print "[ERROR] Third MIN Fit Failed\n%sTotal TS: %s" % (e, self.MIN.logLike.value())
                    print "[INFO] Trying MINUIT algorithm ..."
                    self.initMIN(opt='Minuit', tol=tol2, useBadFit=self.useBadFit)
                    self.MIN.fit(covar=True, optObject=self.MINobj)
                    if ('isRestoreEdisp' in dir(self)) and (not self.isRestoreEdisp): self.restoreEdisp()
                    print '[INFO] MINUIT Fit Finished.'
                    print '[INFO] Fit Quality:  ', self.MINobj.getQuality()
                    print '[INFO] Return Code:  ', self.MINobj.getRetCode()
                    print '[INFO] Fit Distance: ', self.MINobj.getDistance()
                    if self.MINobj.getQuality() < 3:
                        print "[ERROR] MINUIT algorithm failed, bailing"
                        self.modelPass2BAD = '.'.join(modelOut.split('.')[:-1]) + '_BAD.xml'
                        self.MIN.logLike.writeXml(self.modelPass2BAD)
                        print "[INFO] Saved MIN as "+self.modelPass2BAD
                        return False
                    else:
                        self.modelPass2 = modelOut
                        self.MIN.logLike.writeXml(self.modelPass2)
                        print "[INFO] Saved MIN as "+self.modelPass2
                        return True

        print '[INFO] NEWMINUIT Fit Finished.'
        print '[INFO] Total TS: ', self.MIN.logLike.value()
      # print '[INFO] Fit Quality:  ', self.MINobj.getQuality()
        print '[INFO] Return Code:  ', self.MINobj.getRetCode()
        print '[INFO] Fit Distance: ', self.MINobj.getDistance()

        self.fitbit = True
        if self.MINobj.getRetCode() > 0:
            print '[ERROR] NEWMINUIT DID NOT CONVERGE!!!'
            failure = ""
            retCode = self.MINobj.getRetCode() - 100
            if retCode & 1:
                failure += " IsAboveMaxEdm"
            if retCode & 2:
                failure += " HasCovariance"
            if retCode & 4:
                failure += " HesseFailed"
            if retCode & 8:
                failure += " HasMadePosDefCovar"
            if retCode & 16:
                failure += " HasPosDefCovar"
            if retCode & 32:
                failure += " HasAccurateCovar"
            if retCode & 64:
                failure += " HasValidCovariance"
            if retCode & 128:
                failure += " HasValidParameters"
            if retCode & 256:
                failure += " IsValid"
            print '[INFO] ' + failure
            self.modelPass2BAD = '.'.join(modelOut.split('.')[:-1]) + '_BAD.xml'
            self.MIN.logLike.writeXml(self.modelPass2BAD)
            print "[INFO] Saved MIN as "+self.modelPass2BAD
            return False
        else:
            self.modelPass2 = modelOut
            self.MIN.logLike.writeXml(self.modelPass2)
            print "[INFO] Saved MIN as "+self.modelPass2
            return True

    def paramsAtLimit(self, limit=0.1):
        """
        This function will print out any sources whose parameters
        are close to their limits.  You could use this to find sources
        that are having issues being fit.  This function is useful
        when you're having trouble getting convergence from the New
        Minuit fit routine. The limit is in percentage difference of a
        bound.  If one of the bounds is zero it uses the value of the
        parameter to check for closeness (absolute instead of percent
        differenct).  The default is 0.1 (10%) difference for a measure
        of closeness.
        """
        try:
            self.MIN
        except AttributeError:
            raise RuntimeError("MIN object does not exist. \
                                Create it first with the initMIN function and then fit it with the fitMIN function.")

        if not self.fitbit:
            print "[WARN] Fit isn't current, these values might not be correct. Run fitMIN first."

        for src in self.MIN.sourceNames():
            for name in self.MIN[src].funcs['Spectrum'].paramNames:
                bounds = self.MIN[src].funcs['Spectrum'].getParam(name).getBounds()
                value = self.MIN[src].funcs['Spectrum'].getParam(name).value()

                try:
                    distToLower = abs((value - bounds[0])/bounds[0])
                except ZeroDivisionError:
                    distToLower = abs(value)

                try:
                    distToUpper = abs((value - bounds[1])/bounds[1])
                except ZeroDivisionError:
                    distToUpper = abs(value)

                if distToLower < limit:
                    print ('[INFO] The %s (%s) of %s is close (%s) to its lower limit (%s)' %
                           (name, value, src, distToLower, bounds[0]))

                if distToUpper < limit:
                    print ('[INFO] The %s (%s) of %s is close (%s) to its upper limit (%s)' %
                           (name, value, src, distToUpper, bounds[1]))

    def writeCountsSpectra(self, outfile='counts_spectra.fits'):
        try:
            self.MIN
        except AttributeError:
            raise RuntimeError("MIN object does not exist. \
                                Create it first with the initMIN function and then fit it with the fitMIN function.")
        self.MIN.writeCountsSpectra(outfile)

    def writeCovar(self, outfile='covariance_pyLike.dat'):
        if self.currentLike.covar_is_current:
            fh = open(outfile, 'w')
            print >>fh, '-'*20, 'free parameters', '-'*20
            prm_lst = []
            for prm in self.currentLike.params():
                if prm.parameter.isFree():
                    freePrm = prm.srcName+'_'+prm.parameter.getName()
                    print >>fh, '%s\t%s' % (len(prm_lst)+1, freePrm)
                    prm_lst.append(freePrm)

            print >>fh, '-'*20, 'covariance matrix', '-'*20
            covar = np.array(self.currentLike.covariance)
            np.savetxt(fh, covar, delimiter='  ', fmt='% .3e')

            print >>fh, '-'*20, 'correlation matrix', '-'*20
            cov_diag = covar.diagonal()
            if (cov_diag < 0.).any():
                print >>fh, '[ERROR] covar.diagonal < 0'
            else:
                sigma = np.sqrt(cov_diag)
                corr = covar / (sigma * sigma.reshape(-1, 1))
                np.savetxt(fh, corr, delimiter='  ', fmt='% .3f')

                print >>fh, '='*20, 'correlation check', '='*20
                triu = np.abs(np.triu(corr, 1))

                print >>fh, '.'*9, 'Very strong correlation (|corr| >= 0.8)', '.'*9
                indics = (triu >= 0.8).nonzero()
                for i, j in zip(indics[0], indics[1]):
                    print >>fh, '%s <-> %s : %s' % (prm_lst[i], prm_lst[j], corr[i, j])

                print >>fh, '.'*9, 'Strong correlation (0.6 <= |corr| < 0.8)', '.'*8
                indics = ((triu < 0.8)*(triu >= 0.6)).nonzero()
                for i, j in zip(indics[0], indics[1]):
                    print >>fh, '%s <-> %s : %s' % (prm_lst[i], prm_lst[j], corr[i, j])
            fh.close()
        else:
            print '[WARN] No covariance is available!'

    def initSed(self, saveFolder='./', writeModel=True):
        if not saveFolder:
            self.sedSaveFolder = ''
        elif saveFolder.endswith('/'):
            self.sedSaveFolder = saveFolder
        else:
            self.sedSaveFolder = saveFolder + '/'

        self.initDRM()
        self.sedEnergy = self.DRM.energies
        self.saveEdisp()
        if writeModel:
            for src in self.DRM.sourceNames():
                if self.DRM[src].src.fixedSpectrum(): continue
                for src_par in self.DRM[src].funcs['Spectrum'].paramNames:
                    if (src_par == 'Prefactor' or src_par == 'norm' or
                        src_par == 'Integral' or src_par == 'Value' or
                        src_par == 'Normalization'):
                        continue
                    src_index = self.DRM.par_index(src, src_par)
                    self.DRM.freeze(src_index)
            self.srcModel = self.sedSaveFolder + 'model_sed.xml'
            self.DRM.logLike.writeXml(self.srcModel)
        del self.DRM
        self.eBins = []

    def runSed(self, centerSrc, tsLimit=2., tol1=0.001, tol2=1.e-5, ibin_min=None, ibin_max=None):
        Emins, Emaxs = self.sedEnergy[:-1], self.sedEnergy[1:]
        if isinstance(ibin_min, int):
            minBin = max(0, ibin_min)
        else:
            minBin = 0

        if isinstance(ibin_max, int):
            maxBin = min(len(Emins)-1, ibin_max)
        else:
            maxBin = len(Emins)
        print '[INFO] Calculate from bin %s to bin %s ...' % (minBin, maxBin)

        for iband in xrange(minBin, maxBin+1):
            self.initDRM(tol=tol1)
            if not centerSrc in self.DRM.sourceNames():
                raise IOError('%s is not in the model file!' % centerSrc)
            emin, emax = Emins[iband], Emaxs[iband]
            self.setERange(emin, emax)
            self.fitDRM(modelOut=self.sedSaveFolder+'model_pass1_band%s.xml' % iband, pokeSrc=centerSrc)
            self.DRMpristine.restore()

            self.initMIN(tol=tol2)
            self.setERange(emin, emax)
            self.restoreEdisp()
            self.fitMIN(modelOut=self.sedSaveFolder+'model_pass2_band%s.xml' % iband)
            self.paramsAtLimit()

            self.writeResults(self.sedSaveFolder+'results_pyLike_band%s.dat' % iband, centerSrc, tsLimit)
            self.eBins.append(iband)
            del self.MIN, self.MINobj # It MUST be removed to avoid bugs!

    def collectSed(self, centerSrc, output='mysed.dat'):
        self.eBins = sorted(list(set(self.eBins)))
        fh = open(self.sedSaveFolder+output, 'w')
        print >>fh, '#I\tEmin[MeV]\tEmax[MeV]\tflux\tfluxErr\tnpred\tts\tfitQual\tretCode\tEDM\tUL'
        for iband in self.eBins:
            try:
                datafh = open(self.sedSaveFolder+'results_pyLike_band%s.dat' % iband)
            except IOError:
                continue
            info = eval(datafh.read())
            centerSrcInfo = info[centerSrc]
            fitQual = info['zFitQuality']
            retCode = info['zReturnCode']
            EDM = info['zEDM']
            datafh.close()

            emin, emax = self.sedEnergy[iband], self.sedEnergy[iband+1]
            ts = float(centerSrcInfo['TS value'])
            tmp = centerSrcInfo['Flux'].split('+/-')
            flux_ = float(tmp[0])
            fluxErr_ = float(tmp[1])
            npred = float(centerSrcInfo['NPred'])

            try:
                flux_ul = float(centerSrcInfo['Upper Limit'])
                print >> fh, iband, emin, emax, flux_, fluxErr_, npred, ts, fitQual, retCode, EDM, flux_ul
            except KeyError:
                print >> fh, iband, emin, emax, flux_, fluxErr_, npred, ts, fitQual, retCode, EDM
        fh.close()
